package org.acghub.langchain4j.message;

import java.util.List;
import java.util.Map;
import java.util.function.Function;

public final class DefaultMessageMemory<K, V> extends AbstractMessageMemory<K, V> {

    /**
     * The maximum number of messages to keep in the history.
     * 0: no history, >0: max number of messages to keep
     */
    private final int maxHistorySize;

    public DefaultMessageMemory(int maxHistorySize) {
        super();
        this.maxHistorySize = maxHistorySize;
    }

    public DefaultMessageMemory(Map<K, List<V>> memory, Function<K, List<V>> valueListSupplier, int maxHistorySize) {
        super(memory, valueListSupplier);
        this.maxHistorySize = maxHistorySize;
    }

    private void tryTrim(K k) {
        List<V> history = getHistory(k);
        while (maxHistorySize > 0 && history.size() > maxHistorySize) {
            removeFirst(k);
        }
    }

    @Override
    public void add(K k, V v, int index) {
        super.add(k, v, index);
        tryTrim(k);
    }

    @Override
    public void addFirst(K k, V v) {
        super.addFirst(k, v);
    }

    @Override
    public void addLast(K k, V v) {
        super.addLast(k, v);
        tryTrim(k);
    }
}
